import os
import subprocess
import argparse
import re

def detect_gpu_type():
    """检测系统可用的GPU类型"""
    try:
        # 检测NVIDIA GPU
        subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return 'nvidia'
    except (subprocess.CalledProcessError, FileNotFoundError):
        try:
            # 检测AMD GPU
            subprocess.run(['rocminfo'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
            return 'amd'
        except (subprocess.CalledProcessError, FileNotFoundError):
            try:
                # 检测Intel GPU
                subprocess.run(['intel_gpu_top'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
                return 'intel'
            except (subprocess.CalledProcessError, FileNotFoundError):
                return None

def get_gpu_encoder(gpu_type):
    """根据GPU类型返回最佳H.265编码器"""
    if gpu_type == 'nvidia':
        return 'hevc_nvenc'  # NVIDIA H.265编码器
    elif gpu_type == 'amd':
        return 'hevc_amf'    # AMD H.265编码器
    elif gpu_type == 'intel':
        return 'hevc_qsv'    # Intel H.265编码器
    else:
        return 'libx265'     # 回退到CPU H.265编码

def compress_mp4(input_file, output_file=None, overwrite=False, gpu_accel=True, quality='medium'):
    """
    使用H.265编码压缩MP4视频文件（视频重新编码压缩，音频直接拷贝）
    
    参数:
    input_file: 输入的MP4文件路径
    output_file: 输出的MP4文件路径（默认为输入文件同目录同名）
    overwrite: 是否覆盖已存在的输出文件
    gpu_accel: 是否使用GPU加速
    quality: 压缩质量设置（'low', 'medium', 'high'）
    """
    # 设置默认输出文件名
    if output_file is None:
        filename, ext = os.path.splitext(input_file)
        output_file = f"{filename}_compressed{ext}"
    
    # 检查输入文件是否存在
    if not os.path.isfile(input_file):
        raise FileNotFoundError(f"输入文件不存在: {input_file}")
    
    # 处理已存在的输出文件
    if os.path.isfile(output_file):
        if overwrite:
            os.remove(output_file)
        else:
            print(f"错误: 输出文件已存在 - {output_file} (使用 -y 选项覆盖)")
            return False
    
    # 检测GPU类型
    gpu_type = detect_gpu_type() if gpu_accel else None
    encoder = get_gpu_encoder(gpu_type) if gpu_accel else 'libx265'
    
    print(f"使用H.265编码器: {encoder} ({gpu_type or 'CPU'})")
    print(f"质量设置: {quality}")
    
    # 构建FFmpeg命令
    cmd = [
        'ffmpeg',
        '-i', input_file,
        '-c:v', encoder,      # 使用H.265视频编码器
        '-c:a', 'copy',       # 直接拷贝音频流
        '-movflags', '+faststart',  # 优化网络播放
        '-tag:v', 'hvc1',     # 添加H.265标签确保兼容性
        '-y' if overwrite else '-n',  # 自动覆盖/不覆盖
    ]
    
    # 根据质量设置添加参数
    if quality == 'high':
        # 高质量（文件较大）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p7', '-tune', 'hq', '-cq', '18'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'quality', '-qp_i', '18', '-qp_p', '18'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'slower', '-global_quality', '18'])
        else:  # CPU编码
            cmd.extend(['-preset', 'slow', '-crf', '18'])
    elif quality == 'medium':
        # 中等质量（平衡）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p7', '-cq', '23'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'balanced', '-qp_i', '23', '-qp_p', '23'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'medium', '-global_quality', '23'])
        else:  # CPU编码
            cmd.extend(['-preset', 'medium', '-crf', '23'])
    else:  # low
        # 低质量（文件较小）
        if encoder == 'hevc_nvenc':
            cmd.extend(['-preset', 'p7', '-cq', '28'])
        elif encoder == 'hevc_amf':
            cmd.extend(['-quality', 'speed', '-qp_i', '28', '-qp_p', '28'])
        elif encoder == 'hevc_qsv':
            cmd.extend(['-preset', 'fast', '-global_quality', '28'])
        else:  # CPU编码
            cmd.extend(['-preset', 'fast', '-crf', '28'])
    
    cmd.append(output_file)
    
    try:
        # 执行转换命令
        print("开始压缩...")
        process = subprocess.Popen(
            cmd,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            encoding='utf-8',
            errors='replace'
        )
        
        # 实时显示进度
        while True:
            line = process.stderr.readline()
            if not line:
                break
            # 提取并显示进度信息
            if 'time=' in line:
                time_match = re.search(r'time=(\d+:\d+:\d+\.\d+)', line)
                if time_match:
                    print(f"\r压缩进度: {time_match.group(1)}", end='', flush=True)
        
        # 等待进程完成
        process.wait()
        
        if process.returncode == 0:
            print(f"\n压缩成功: {input_file} -> {output_file}")
            return True
        else:
            print(f"\n压缩失败，返回码: {process.returncode}")
            # 打印错误信息
            print("错误信息:")
            # 读取剩余的错误输出
            while True:
                line = process.stderr.readline()
                if not line:
                    break
                print(line, end='')
            return False
    except FileNotFoundError:
        print("\n错误: 未找到FFmpeg，请先安装FFmpeg并添加到系统路径")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='使用GPU加速压缩MP4视频文件（使用H.265编码，视频重新编码，音频直接拷贝）')
    parser.add_argument('input', help='输入MP4文件路径')
    parser.add_argument('-o', '--output', help='输出MP4文件路径（可选）')
    parser.add_argument('-y', '--overwrite', action='store_true', help='覆盖已存在的输出文件')
    parser.add_argument('--no-gpu', action='store_true', help='禁用GPU加速（使用CPU编码）')
    parser.add_argument('-q', '--quality', choices=['low', 'medium', 'high'], default='medium',
                        help='压缩质量设置（low, medium, high），默认为medium')
    
    args = parser.parse_args()
    
    compress_mp4(
        input_file=args.input,
        output_file=args.output,
        overwrite=args.overwrite,
        gpu_accel=not args.no_gpu,
        quality=args.quality
    )